#ifndef SHELL_ORDERING_H_
#define SHELL_ORDERING_H_
/*
 * shell_ordering.h
 *
 *  Created on: Mar 29, 2013
 *      Author: drewlewis
 */

#include <chemistry/molecule/molecule.h>
#include "../Global.h"
#include <Eigen/Dense>
#include <Eigen/StdVector>
#include "kcenter.h"
#include <iostream>
#include <cmath>
#include <utility>

class AtomOrdering {

    private:
        typedef  sc::Ref<sc::Molecule> Mol;
        typedef  Eigen::aligned_allocator<Vector3> EigenStd;
        typedef  std::pair<int, Vector3> pair_int;

        class Atom_{
                //Class to hold an atom since tuple appears to be hard to get to work
            public:
                int charge_;
                double mass_;
                Vector3 pos_;

                Atom_(int charge, double mass, Vector3 pos) :
                    charge_(charge),
                    mass_(mass),
                    pos_(pos){};
        }; //Atom_

        struct norm_sort{
             /* This is a bit complicated so lets see if I can spell it out.
              * First compute the difference between the
              * two atoms and the position
              * that we gave to the object constructor.
              * This can be any Eigen::Vector3d.
              * We can chose any atom as the first atom in the molecule by
              * passing in its position as the pos_ vector.
              *
              * The long return things basically says the following:
              *     if(diff1 < diff2)
              *         true;
              *     else if(diff1 == diff2 && atom1.pos_[0] < atom2.pos_[0])
              *         true;
              *     else if(diff1 == diff2 && atom1.pos_[0] == atom2.pos_[0] &&
              *         atom1.pos_[1] < atom2.pos_[1])
              *         true;
              *     else if (diff1 == diff2 && atom1.pos_[0] == atom2.pos_[0] &&
              *         atom1.pos_[1] == atom2.pos_[1] &&
              *         atom1.pos_[2] < atom.pos_[2])
              *         true;
              *     else false
              */
             bool operator() (Atom_ atom1, Atom_ atom2) const {

                 // sort by distance from input atom
                 Vector3 pos1 = atom1.pos_;
                 Vector3 pos2 = atom2.pos_;
                 double norm1 = (pos1 - pos_).norm();
                 double norm2 = (pos2 - pos_).norm();
                 return( (norm1 < norm2)  ||
                         ( norm1 == norm2 && pos1[0] < pos2[0] ) ||
                         ( norm1 == norm2 && pos1[0] == pos2[0]
                                          && pos1[1] < pos2[1] ) ||
                         ( norm1 == norm2 && pos1[0] == pos2[0]
                                && pos1[1] == pos2[1] && pos1[2] < pos2[2] ) );
             }

             norm_sort(Vector3 pos) : pos_(pos) {}

             private:
                 Vector3 pos_;

        };

        struct IsHeavy{

            IsHeavy(double mass) : heavy_mass_(mass) {};

            bool operator () (Atom_ atom) const {
                return (atom.mass_ == heavy_mass_);
            }

            private:
                double heavy_mass_;
        };

        struct find_heaviest{
            bool operator () (Atom_ atom1, Atom_ atom2) const {
                return (atom1.mass_ < atom2.mass_);
            }
        } f_heavy;

        void
        canaonical_sort(){

            Atom_ heaviest_atom = *std::max_element(atoms_.begin(), atoms_.end(),
                                               f_heavy);
            double heaviest = heaviest_atom.mass_;
            std::cout << "The heaviest Atom has mass = " << heaviest
                      << std::endl;

            // Figure out number of heaviest atoms
            int nheavy = std::count_if(atoms_.begin(), atoms_.end(),
                                       IsHeavy(heaviest));

            Vector3 pos = Eigen::VectorXd::Zero(3);

            std::sort(atoms_.begin(), atoms_.end(), norm_sort(pos));
        }

        void
        guess_kcenters(size_t nkcenters){
            // This function will determine which atoms to make center guesses
            // for the k-mean's clustering algorithm.

            nkcenters_ = nkcenters;

            // Sort atoms by
            // 1. Find heaviest atom
            // 2. Find farthest from origin heaviest atom
            // 3. Sort other atoms based on distance from atom in step 2. :w
            canaonical_sort();

            // Make kcenter guess = the first n heaviest atoms.
            for(std::size_t i = 0; i < nkcenters; ++i){
                kcenters_.push_back( KCenter( atoms_[i].pos_) );
            }
        }

        size_t find_nearest_kcenter(Vector3 &vec){
            //Function will determine the closes kcenter by computing the
            // Euclidian norm of the difference vector

            // Just guess that the first kcenter is the cloeset
            Vector3 shortest = kcenters_[0].center() - vec;
            size_t kindex = 0;

            //loop over kcenters
            for(size_t i = 1; i < nkcenters_; ++i){

                //Size of difference vector of atom with each kcenter
                Vector3 diff = kcenters_[i].center() - vec;

                // if we find a closer kcenter switch
                if(diff.norm() < shortest.norm()){
                    kindex = i;
                    shortest = diff;
                }

            }

            return kindex;
        }

        void assign_to_kcenters(){
            // assign every atom to a kcenter

            for(size_t i = 0; i < n_centers_; ++i){

                // Get atom info
                Vector3 pos = atoms_[i].pos_;
                int charge = atoms_[i].charge_;

                // Return index of closest kcenter
                size_t kindex = find_nearest_kcenter(pos);

                // add atom to closes kcenter
                kcenters_[kindex].add_neighbor(charge, pos);

            }

        }



public:
    AtomOrdering(const Mol &mol):mol_(mol), n_centers_(mol->natom()),
        nkcenters_(0), kcenters_() {

        // Give AtomsOrdering Class info from each atom
        for(std::size_t i = 0; i < n_centers_; ++i){ //TODO make parallel later

            // atom info
            double mass = mol->mass(i);
            int charge = mol->charge(i);
            Vector3 pos;

            pos[0] = mol_->r(i,0);
            pos[1] = mol_->r(i,1);
            pos[2] = mol_->r(i,2);

            // push atom object into the list of atoms
            atoms_.push_back( Atom_(charge, mass, pos) );
        }
    }

    void print_mass(){

        //loop over list of atoms
        for(std::size_t i = 0; i < atoms_.size(); ++i){
            std::cout << atoms_[i].mass_ << std::endl;
        }

    }

    void atom_clustering(const std::size_t ncenters){
        // Will determing the best grouping of atoms for tiles based on
        // minimizing the distance between atoms.

        // Assigns kcenters to the first n heavest atoms
        if(ncenters < n_centers_)
            guess_kcenters(ncenters);
        else // Never have more centers than the number of atoms.
            guess_kcenters(n_centers_);

        // Puts every atom into its closest center if equidistant then which ever
        // center was listed first
        assign_to_kcenters();

        // Number of Lloyd's algorithm iterations
        for(size_t i = 0; i < 10; ++i){

            // Loop over kcenters
            for(size_t j = 0; j < nkcenters_; ++j){

                // calculate new kcenter based on objects stored in center
                KCenter center(kcenters_[j].centroid() );

                // Get rid of old center data
                kcenters_[j].remove_data();

                //create a new center
                kcenters_[j] = center;
            }

            // Now that new centers have been created put atoms into closest kcenter
            assign_to_kcenters();
        }

    }

    void print_kcenters(){
        // Loop over kcenters
        for(std::size_t i = 0; i < nkcenters_; ++i ){
            std::cout << "Kcenter " << i << " has center " << std::endl;
            kcenters_[i].print_center();
            std::cout << std::endl;
        }
    }

    void print_kcenters_xyz(){
        // Loop over kcenters

        std::map<int, std::string> Nucli;
        Nucli[6] = "C";
        Nucli[1] = "H";
        Nucli[8] = "O";
        Nucli[7] = "N";
        Nucli[16]= "S";
        Nucli[26]= "Fe";
        Nucli[20]= "Ca";

        for(std::size_t i = 0; i < nkcenters_; ++i ){
            std::size_t n_neighbors = kcenters_[i].n_neighbors();

            std::cout << n_neighbors << "\n"
                            "Group " << i << std::endl;

            for(std::size_t j = 0; j < n_neighbors; ++j){
                printf("  %s \t %f \t %f \t %f \n",
                       Nucli[kcenters_[i].neighbor_charge(j)].c_str(),
                       kcenters_[i].neighbor_pos(j,0),
                       kcenters_[i].neighbor_pos(j,1),
                       kcenters_[i].neighbor_pos(j,2));
            }
            std::cout << std::endl;
        }
    }

    void write_kcenters_xyz(){
        std::string out = "fragment";

        // Loop over kcenters

        std::map<int, std::string> Nucli;
        Nucli[6] = "C";
        Nucli[1] = "H";
        Nucli[7] = "N";
        Nucli[8] = "O";
        Nucli[16] = "S";
        Nucli[15] = "P";
        Nucli[26]= "Fe";
        Nucli[20]= "Ca";
        Nucli[30]= "Zn";

        for(std::size_t i = 0; i < nkcenters_; ++i ){
            std::string filename;
            std::ostringstream convert;
            convert << i;
            filename = convert.str();
            std::string xyz(".xyz");
            filename = out + filename + xyz;
            std::ofstream myfile;
            myfile.open(filename.c_str());


            std::size_t n_neighbors = kcenters_[i].n_neighbors();

            myfile << n_neighbors << "\n"
                            "Group " << i << std::endl;

            for(std::size_t j = 0; j < n_neighbors; ++j){
                myfile << Nucli[kcenters_[i].neighbor_charge(j)] << "\t" <<
                                kcenters_[i].neighbor_pos(j,0) << "\t" <<
                                kcenters_[i].neighbor_pos(j,1) << "\t" <<
                                kcenters_[i].neighbor_pos(j,2) << std::endl;
            }
            myfile << std::endl;
            myfile.close();
        }

        std::string molecule("molecule.xyz");
        std::ofstream myfile;
        myfile.open(molecule.c_str());
        std::cout << atoms_.size() << "\n" << std::endl;
        myfile << atoms_.size() << "\nWhole Molecule" <<  std::endl;
        for(std::size_t i = 0; i < nkcenters_; ++i ){
            std::size_t n_neighbors = kcenters_[i].n_neighbors();
            for(std::size_t j = 0; j < n_neighbors; ++j){
                myfile << Nucli[kcenters_[i].neighbor_charge(j)] << "\t" <<
                                kcenters_[i].neighbor_pos(j,0) << "\t" <<
                                kcenters_[i].neighbor_pos(j,1) << "\t" <<
                                kcenters_[i].neighbor_pos(j,2) << std::endl;
            }
        }
        myfile << std::endl;
        myfile.close();
    }

    Mol get_optimized_mol(){
        // Return a molecule object with the ordering of the atoms being ordered
        // by kcenter.  First n atoms belong to center 0, next m atoms belong to
        // center 1 and so on.
        Mol mol = new sc::Molecule;

        //Loop over kcenters
        for(size_t i = 0; i < nkcenters_; ++i){
            size_t n_neighbors = kcenters_[i].n_neighbors();

            // Loop over atoms in kcenter
            for(size_t j = 0; j < n_neighbors; ++j){
                mol->add_atom(kcenters_[i].neighbor_charge(j),
                              kcenters_[i].neighbor_pos(j,0),
                              kcenters_[i].neighbor_pos(j,1),
                              kcenters_[i].neighbor_pos(j,2) );
            }
        }
        return mol;
    }

    std::vector<std::size_t> get_center_index(){

        std::vector<std::size_t> index;
        index.push_back(0);

        for(std::size_t i = 0; i < nkcenters_; ++i){
            index.push_back(index[i] + kcenters_[i].n_neighbors());
        }

        return index;
    }

private:
    Mol mol_;
    std::size_t n_centers_;
    std::vector< Atom_ > atoms_;
    std::vector< KCenter> kcenters_;
    size_t nkcenters_;
}; // Class ShellOrdering


#endif /* SHELL_ORDERING_H_ */
